##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggpubr)
library(patchwork)

##########################################################################################

option_list <- list(
    make_option(c("--sample_file"), type = "character") ,
    make_option(c("--seg_file"), type = "character") ,
    make_option(c("--class_order_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    sample_file <- paste(work_dir,"/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv",sep="")
    seg_file <- paste(work_dir,"/titan/Titan_all_seg.final.tsv",sep="")
    class_order_file <- paste(work_dir,"/config/Class_order.list",sep="")
    ######
    images_path <- paste(work_dir,"/images/cnv_burden",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_file <- opt$sample_file
seg_file <- opt$seg_file
class_order_file <- opt$class_order_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

info <- data.frame(fread(sample_file))
seg <- data.frame(fread(seg_file))
class_order <- data.frame(fread(class_order_file))

t_logR <- 0.2

###########################################################################################
col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")
col <- col[1:4]

col_im <- brewer.pal(9,"YlGnBu")[6:8]
names(col_im) <- c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)")

###########################################################################################

dat <- seg
dat$Sample <- gsub( "S_" , "S" , dat$Sample )
info$Tumor <- gsub( "S_" , "S" , info$Tumor )
info$Normal <- gsub( "S_" , "S" , info$Normal )

dat$seg.mean <- dat$Median_logR
dat$loc.start <- dat$Start
dat$loc.end <- dat$End
dat$Tumor <- sapply( strsplit( dat$Sample , "_" ) , "[" , 1)
dat$Normal <- sapply( strsplit( dat$Sample , "_" ) , "[" , 2)

dat <- merge( dat , info[,c("Tumor" , "Class" , "Type" , "TCGA_Class")] , by = "Tumor" )
colnames(dat)[ncol(dat)] <- "Molecular.subtype"

###########################################################################################
### 计算CNV改变负荷

dat$TCN <- 2*2^dat$seg.mean
dat$Length <- as.numeric(abs(dat$loc.start - dat$loc.end))

loss_region <- dat[which(dat$seg.mean <= -t_logR),]
gain_region <- dat[which(dat$seg.mean >= t_logR),]
loh_region <- dat[which(dat$MajorCN == 2 & dat$MinorCN == 0),]

result_loss <- c()
result_gain <- c()
result_loh <- c()
for(Tumor in unique(dat$Tumor)){
  print(Tumor)
  class_s <- unique(dat[which(dat$Tumor==Tumor),'Class'])
  normal <- unique(dat[which(dat$Tumor==Tumor),'Normal'])
  type <- unique(dat[which(dat$Tumor==Tumor),'Type'])
  Molecular.subtype <- unique(dat[which(dat$Tumor==Tumor),'Molecular.subtype'])
  
  loss_rate <- sum(loss_region[which(loss_region$Tumor==Tumor),'Length'])/sum(dat[which(dat$Tumor==Tumor),'Length'])
  result_loss <- rbind(result_loss,data.frame(Tumor=Tumor,Normal=normal,Rate=loss_rate,Class=class_s,Type= type,Molecular.subtype=Molecular.subtype))

  gain_rate <- sum(gain_region[which(gain_region$Tumor==Tumor),'Length'])/sum(dat[which(dat$Tumor==Tumor),'Length'])
  result_gain <- rbind(result_gain,data.frame(Tumor=Tumor,Normal=normal,Rate=gain_rate,Class=class_s,Type= type,Molecular.subtype=Molecular.subtype))
  
  loh_rate <- sum(loh_region[which(loh_region$Tumor==Tumor),'Length'])/sum(dat[which(dat$Tumor==Tumor),'Length'])
  result_loh <- rbind(result_loh,data.frame(Tumor=Tumor,Normal=normal,Rate=gain_rate,Class=class_s,Type= type,Molecular.subtype=Molecular.subtype))
}

## 确定上下阈值
min_burden <- -0.0001
max_burden <- 0.6

result_gain$Class <- factor(result_gain$Class,levels= unique(class_order$Class), ordered=TRUE)
result_loss$Class <- factor(result_loss$Class,levels= unique(class_order$Class), ordered=TRUE)
result_loh$Class <- factor(result_loh$Class,levels= unique(class_order$Class), ordered=TRUE)

###########################################################################################
## 同一个人IGC和DGC的样本合并
## 分情况画
## IM + IGC
## IM + DGC
## 一个人的多个样本取均值
## loss
dat <- result_loss
dat_loss<- c() 
for(Normal in unique(dat$Normal)){
  print(Normal)
  tmp <- dat[which(dat$Normal==Normal),]

  tmp <- tmp %>%
  group_by( Normal , Class , Type , Molecular.subtype ) %>%
  summarize( CNV_Burden_Loss = median(Rate) )

  tmp <- na.omit(data.frame(tmp))

  dat_loss <- rbind(dat_loss,tmp)
}

## gain
dat <- result_gain
dat_gain <- c() 
for(Normal in unique(dat$Normal)){
  print(Normal)
  tmp <- dat[which(dat$Normal==Normal),]

  tmp <- tmp %>%
  group_by( Normal , Class , Type , Molecular.subtype ) %>%
  summarize( CNV_Burden_Gain = median(Rate) )

  tmp <- na.omit(data.frame(tmp))
  dat_gain <- rbind(dat_gain,tmp)
}

## loh
dat <- result_loh
dat_loh <- c() 
for(Normal in unique(dat$Normal)){
  print(Normal)
  tmp <- dat[which(dat$Normal==Normal),]

  tmp <- tmp %>%
  group_by( Normal , Class , Type , Molecular.subtype ) %>%
  summarize( CNV_Burden_Loh = median(Rate) )

  tmp <- na.omit(data.frame(tmp))
  dat_loh <- rbind(dat_loh,tmp)
}


###########################################################################################
## 不同亚型的肠化和其对应胃癌的关系
trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("P == ",down," %*% 10","^",up)
    return(text)
}

plotFunction <- function(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title){

  dat_tmp <- c()

  for(type in unique(dat_plot$Type) ){

    dat_plot_tmp <- subset( dat_plot , Type == type )

    a <- dat_plot_tmp[dat_plot_tmp$Class==unique(dat_plot_tmp$Class)[1],"Burden_use"]
    b <- dat_plot_tmp[dat_plot_tmp$Class==unique(dat_plot_tmp$Class)[2],"Burden_use"]
    p <- wilcox.test( a , b )$p.value

    if( p < 0.01 ){
        p_text <- trans(p)
    }else{
        p_text <- paste0( "P == " , round(as.numeric(p) , 3) ) 
    }
    dat_plot_tmp$p_text <- ""
    dat_plot_tmp$p_text[1] <- p_text
    dat_tmp <- rbind( dat_plot_tmp , dat_tmp )
  }
  col_tmp <- c(
    rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
    rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
    rgb(red=2,green=100,blue=190,alpha=255,max=255) 
  )

  plot <- ggplot(data=dat_tmp , mapping = aes(x=Class,y=Burden_use))+
  geom_line( aes( group = Normal ) , size = 0.3 , color = "gray" ) +
  geom_boxplot(lwd=1.5,aes(color=Class) , outlier.shape = NA , size = 0.2) +
  geom_jitter(position=position_jitter(0.2),aes(color=Class) , size = 1) +
  scale_color_manual(values=col_tmp) +
  facet_grid(.~Type , scales="free_x" ) +
  xlab(NULL) +
  ylab(y_lab)+
  #labs(title=title) +
  theme_bw() +
  #stat_compare_means(comparisons = my_comparisons_1) +
  geom_text(aes(label=p_text , y = 0.6 , x = 1.5),parse = TRUE,size=3 , color = "black") +
  ylim(min_burden,max_burden) +
  theme(
      legend.position = 'none',
      legend.title = element_blank() ,
      panel.grid.major=element_blank(),
      panel.grid.minor=element_blank(),
      panel.background = element_blank(),
      panel.border = element_blank(),
      plot.title = element_text(size = 12,color="black",face='bold'),
      legend.text = element_text(size = 12,color="black",face='bold'),
      axis.text.y = element_text(size = 12,color="black",face='bold'),
      axis.title.x = element_text(size = 12,color="black",face='bold'),
      axis.title.y = element_text(size = 12,color="black",face='bold'),
      axis.text.x = element_text(size = 12,color="black",face='bold') ,
      axis.ticks.length = unit(0.2, "cm") ,
      strip.text.x = element_text(size = 14, colour = "black",face='bold') ,
      axis.line = element_line(size = 0.5)) 
  return(plot)
}

my_comparisons_1 <- list( c(1, 2) )
for( molecular in c("All" , "GS" , "CIN" , "MSI") ){

  if(molecular == "All"){
    dat_plot_tmp_gain <- dat_gain
    dat_plot_tmp_loss <- dat_loss
    dat_plot_tmp_loh <- dat_loh
  }else{
    dat_plot_tmp_gain <- subset( dat_gain , Molecular.subtype == molecular )
    dat_plot_tmp_gain <- subset( dat_gain , Normal %in% dat_plot_tmp_gain$Normal )

    dat_plot_tmp_loss <- subset( dat_loss , Molecular.subtype == molecular )
    dat_plot_tmp_loss <- subset( dat_loss , Normal %in% dat_plot_tmp_loss$Normal )

    dat_plot_tmp_loh <- subset( dat_loh , Molecular.subtype == molecular )
    dat_plot_tmp_loh <- subset( dat_loh , Normal %in% dat_plot_tmp_loss$Normal )
  }

  dat_plot_tmp_gain <- subset( dat_plot_tmp_gain , Type != "IM + IGC + DGC" )
  dat_plot_tmp_loss <- subset( dat_plot_tmp_loss , Type != "IM + IGC + DGC" )
  dat_plot_tmp_loh <- subset( dat_plot_tmp_loh , Type != "IM + IGC + DGC" )

  title <- molecular

  ## 扩增负荷
  y_lab <- "Fraction of CNV altered\n(gain)"
  dat_plot <- dat_plot_tmp_gain
  dat_plot$Burden_use <- dat_plot$CNV_Burden_Gain
  dat_plot <- dat_plot %>%
  group_by( Normal , Class , Type ) %>%
  summarize( Burden_use = median(Burden_use) )
  dat_plot$Class <- factor(dat_plot$Class,levels=c( "IM" , "IGC" , "DGC") , ordered=T)
  ## 标记样本数量
  sample_num <- dat_plot %>%
  group_by( Type ) %>%
  summarize( nums = length(unique(Normal)) )
  dat_plot <- merge( dat_plot , sample_num , by = "Type" )
  dat_plot$Type <- factor( dat_plot$Type , levels = c("IM + IGC" , "IM + DGC") , order = T )
  #dat_plot$Type <- paste0( dat_plot$Type , "(" , dat_plot$nums , ")" )
  p1 <- plotFunction(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title)
  
  ## 丢失负荷
  y_lab <- "Fraction of CNV altered\n(loss)"
  dat_plot <- dat_plot_tmp_loss
  dat_plot$Burden_use <- dat_plot$CNV_Burden_Loss
  dat_plot <- dat_plot %>%
  group_by( Normal , Class , Type  ) %>%
  summarize( Burden_use = median(Burden_use) )
  dat_plot$Class <- factor(dat_plot$Class,levels=c( "IM" , "IGC" , "DGC") , ordered=T)
  sample_num <- dat_plot %>%
  group_by( Type ) %>%
  summarize( nums = length(unique(Normal)) )
  dat_plot <- merge( dat_plot , sample_num , by = "Type" )
  #dat_plot$Type <- paste0( dat_plot$Type , "(" , dat_plot$nums , ")" )
  p2 <- plotFunction(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title)

  ## LOH负荷
  y_lab <- "Fraction of CNV altered\n(loh)"
  dat_plot <- dat_plot_tmp_loh
  dat_plot$Burden_use <- dat_plot$CNV_Burden_Loh
  dat_plot <- dat_plot %>%
  group_by( Normal , Class , Type  ) %>%
  summarize( Burden_use = median(Burden_use) )
  dat_plot$Class <- factor(dat_plot$Class,levels=c( "IM" , "IGC" , "DGC") , ordered=T)
  sample_num <- dat_plot %>%
  group_by( Type ) %>%
  summarize( nums = length(unique(Normal)) )
  dat_plot <- merge( dat_plot , sample_num , by = "Type" )
  #dat_plot$Type <- paste0( dat_plot$Type , "(" , dat_plot$nums , ")" )
  p3 <- plotFunction(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title)


  ## 合并
  images_name <- paste(images_path,"/CNV_burden.IM_GC.",molecular,".pdf",sep="")
  result_p <- p1 / p2
  ggsave(file=images_name,plot=result_p,width=3.1/0.8,height=4.8/0.8)


  result <- dat_plot_tmp_gain %>%
  group_by( Class , Type  ) %>%
  summarize( Burden_use = median(CNV_Burden_Gain) )
  images_name <- paste(images_path,"/CNV_burden.IM_GC.",molecular,".gain.tsv",sep="")
  write.table(result , images_name , row.names = F , sep = "\t" , quote = F)

  result <- dat_plot_tmp_loss %>%
  group_by( Class , Type  ) %>%
  summarize( Burden_use = median(CNV_Burden_Loss) ) 
  images_name <- paste(images_path,"/CNV_burden.IM_GC.",molecular,".loss.tsv",sep="")
  write.table(result , images_name , row.names = F , sep = "\t" , quote = F)

  result <- dat_plot_tmp_loh %>%
  group_by( Class , Type  ) %>%
  summarize( Burden_use = median(CNV_Burden_Loh) ) 
  images_name <- paste(images_path,"/CNV_burden.IM_GC.",molecular,".loh.tsv",sep="")
  write.table(result , images_name , row.names = F , sep = "\t" , quote = F)

  images_name <- paste(images_path,"/CNV_burden.IM_GC.",molecular,".loh.pdf",sep="")
  result_p <- p3
  ggsave(file=images_name,plot=result_p,width=3.1,height=1.8)

}

###########################################################################################
## 不同胃癌亚型CNV负荷比较
max_burden <- 0.2
plotFunction <- function(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title){
  plot <- ggplot(data=dat_plot,mapping = aes(x=Class,y=Burden_use))+
  geom_boxplot(lwd=1.5,aes(color=Class) , outlier.shape = NA) +
  geom_jitter(position=position_jitter(0.1),aes(color=Class)) +
  scale_color_manual(values=as.character(col_im)) +
  facet_grid(.~Type , scales="free_x" ) +
  xlab(NULL) +
  ylab(y_lab)+
  labs(title=title) +
  theme_bw() +
  stat_compare_means(comparisons = my_comparisons_1 , label.y = c(0.05, 0.07, 0.09) , method = "wilcox.test") +
  ylim(min_burden,max_burden) +
  theme(
      legend.position = 'none',
      legend.title = element_blank() ,
      panel.grid.major=element_blank(),
      panel.grid.minor=element_blank(),
      panel.background = element_blank(),
      panel.border = element_blank(),
      plot.title = element_text(size = 12,color="black",face='bold'),
      legend.text = element_text(size = 12,color="black",face='bold'),
      axis.text.y = element_text(size = 12,color="black",face='bold'),
      axis.title.x = element_text(size = 12,color="black",face='bold'),
      axis.title.y = element_text(size = 12,color="black",face='bold'),
      axis.text.x = element_text(size = 12,color="black",face='bold') ,
      axis.ticks.length = unit(0.2, "cm") ,
      strip.text.x = element_text(size = 13, colour = "black",face='bold') ,
      axis.line = element_line(size = 0.5)) 
}

cin_sample_use <- unique(c(subset(dat_gain , Molecular.subtype=="CIN")$Normal , subset(dat_loss , Molecular.subtype=="CIN")$Normal))
gs_sample_use <- unique(c(subset(dat_gain , Molecular.subtype=="GS")$Normal , subset(dat_loss , Molecular.subtype=="GS")$Normal))
msi_sample_use <- unique(c(subset(dat_gain , Molecular.subtype=="MSI")$Normal , subset(dat_loss , Molecular.subtype=="MSI")$Normal))
my_comparisons_1 <- list( c(1, 2) , c(1,3) , c(2,3) )

for( type in c("All" , "IM + IGC" , "IM + DGC") ){

  if(type == "All"){
    dat_plot_tmp_gain <- dat_gain
    dat_plot_tmp_loss <- dat_loss
  }else{
    dat_plot_tmp_gain <- subset( dat_gain , Type == type )
    dat_plot_tmp_loss <- subset( dat_loss , Type == type )
  }

  dat_plot_tmp_gain <- subset( dat_plot_tmp_gain , Class == "IM" & Type != "IM + IGC + DGC" )
  dat_plot_tmp_loss <- subset( dat_plot_tmp_loss , Class == "IM" & Type != "IM + IGC + DGC" )

  dat_plot_tmp_gain$Class <- ""
  dat_plot_tmp_gain$Class <- ifelse( dat_plot_tmp_gain$Normal %in% gs_sample_use , "GS" , dat_plot_tmp_gain$Class )
  dat_plot_tmp_gain$Class <- ifelse( dat_plot_tmp_gain$Normal %in% cin_sample_use , "CIN" , dat_plot_tmp_gain$Class )
  dat_plot_tmp_gain$Class <- ifelse( dat_plot_tmp_gain$Normal %in% msi_sample_use , "MSI" , dat_plot_tmp_gain$Class )
  dat_plot_tmp_gain <- subset( dat_plot_tmp_gain , Class != "" )

  dat_plot_tmp_loss$Class <- ""
  dat_plot_tmp_loss$Class <- ifelse( dat_plot_tmp_loss$Normal %in% gs_sample_use , "GS" , dat_plot_tmp_loss$Class )
  dat_plot_tmp_loss$Class <- ifelse( dat_plot_tmp_loss$Normal %in% cin_sample_use , "CIN" , dat_plot_tmp_loss$Class )
  dat_plot_tmp_loss$Class <- ifelse( dat_plot_tmp_loss$Normal %in% msi_sample_use , "MSI" , dat_plot_tmp_loss$Class )
  dat_plot_tmp_loss <- subset( dat_plot_tmp_loss , Class != "" )

  title <- type

  ## 扩增负荷
  y_lab <- "Fraction of CNV altered\n(gain)"
  dat_plot <- dat_plot_tmp_gain
  dat_plot$Burden_use <- dat_plot$CNV_Burden_Gain
  dat_plot <- dat_plot %>%
  group_by( Normal , Class , Type ) %>%
  summarize( Burden_use = median(Burden_use) )
  sample_num <- dat_plot %>%
  group_by( Class ) %>%
  summarize( nums = length(unique(Normal)) )
  dat_plot <- merge( dat_plot , sample_num , by.x = "Class" )
  dat_plot$Class <- paste0( dat_plot$Class , "\n(" , dat_plot$nums , ")" )
  p1 <- plotFunction(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title)
  
  ## 丢失负荷
  y_lab <- "Fraction of CNV altered\n(loss)"
  dat_plot <- dat_plot_tmp_loss
  dat_plot$Burden_use <- dat_plot$CNV_Burden_Loss
  dat_plot <- dat_plot %>%
  group_by( Normal , Class , Type  ) %>%
  summarize( Burden_use = median(Burden_use) )
  sample_num <- dat_plot %>%
  group_by( Class ) %>%
  summarize( nums = length(unique(Normal)) )
  dat_plot <- merge( dat_plot , sample_num , by.x = "Class" )
  dat_plot$Class <- paste0( dat_plot$Class , "\n(" , dat_plot$nums , ")" )
  p2 <- plotFunction(dat_plot = dat_plot , y_lab = y_lab , my_comparisons_1 = my_comparisons_1 , title = title)

  ## 合并
  images_name <- paste(images_path,"/CNV_burden.IM_GC.CNVType.",type,".pdf",sep="")
  result_p <- p1 / p2
  ggsave(file=images_name,plot=result_p,width=4.3/1.2,height=5.8/1.2)



}